# -*- coding: utf-8 -*-
"""
Created on Sat Apr  2 14:38:51 2022

@author: 
"""

import torch
import torch.nn as nn
import numpy as np

def grl_hook(coeff):
    def fun1(grad):
        return -coeff*grad.clone()
    return fun1

class CrossEntropyLabelSmooth(nn.Module):
    def __init__(self, num_classes, epsilon=0.1, use_gpu=True, size_average=True):
        super(CrossEntropyLabelSmooth, self).__init__()
        self.num_classes = num_classes
        self.epsilon = epsilon
        self.use_gpu = False
            #use_gpu
        self.size_average = size_average
        self.logsoftmax = nn.LogSoftmax(dim=1)
        
    def forward(self, inputs, targets):
        log_probs = self.logsoftmax(inputs)
        targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).cpu(), 1)
        if self.use_gpu:
            targets = targets.cuda()
        targets = (1-self.epsilon) * targets + self.epsilon / self.num_classes
        if self.size_average:
            loss = (- targets * log_probs).mean(0).sum()
        else:
            loss = (- targets * log_probs).sum(1)
        return loss

def DANN(input_list, ad_net, entropy=None, coeff=None):
    feature = input_list[0]
    softmax_output = input_list[1]
    ad_out = ad_net(feature)
    # ad_out = nn.Sigmoid()(ad_out)
    
    batch_size = ad_out.size(0) // 2
    dc_target = torch.from_numpy(np.array([[1]] * batch_size + [[0]] * batch_size)).float()
        #.cuda()
    
    if entropy is not None:
        entropy.register_hook(grl_hook(coeff))
        entropy = 1.0 + torch.exp(-entropy)
        
        source_mask = torch.ones_like(entropy)
        source_mask[softmax_output.size(0)//2:] = 0
        source_weight = entropy * source_mask
        
        target_mask = torch.ones_like(entropy)
        target_mask[0:softmax_output.size(0)//2] = 0
        target_weight = entropy * target_mask
        
        weight = source_weight / torch.sum(source_weight).detach().item() + \
            target_weight / torch.sum(target_weight).detach().item()
        return torch.sum(weight.view(-1, 1) * nn.BCELoss(reduce=None)(ad_out, dc_target)) / torch.sum(weight).detach().item()
    else:
        return nn.BCELoss()(ad_out, dc_target)
    
    

                    
        
    
        
    